# 获取字段.py (修改版，合并结果并按1万行分文件输出)
import pymysql
import datetime
from decimal import Decimal
import csv
import sys
import os

# --- 数据库连接配置 ---
host = 'tenant-2102622389-cn-guangzhou-public.bytehouse.volces.com'
port = 3306
database = 'strategy_cloud'
username = 'bytehouse'
password = 'gd3gjGTXTb:WBSEjA67dI'
# ----------------------

# 要查询的表名
table_name = 'cmh_jcy_cmh_std_sales_phl_shard'

# 要获取 distinct 值的字段列表
string_columns = [
    'years', 'quarters', 'months', 'region_type', 'region',
    'sort1', 'sort2', 'sort3', 'sort4', 'tym', 'pm', 'spm',
    'cj', 'rights', 'zx', 'otc_rx', 'jx', 'dx', 'take_way',
    'company_type', 'transfer_rat', 'vpd', 'base_drug', 'health_pro_type',
    'health_type_hist', 'zb_batch_num', 'consistency_evaluation',
    'chemical_drug_cr', 'drug_class', 'short_brand', 'pm_all', 'short_cj',
    'indication_type_name', 'gg1', 'big_jx', 'small_jx', 'first_generics',
    'dsc_import', 'element'
]
string_columns = [
    'brand'
]

# 输出 CSV 文件名前缀和大小限制
output_prefix = 'distinct_values_part_'
rows_per_file = 10000 # 设置为 1 万行

# CSV 文件相关变量
csv_writer = None
current_file = None
current_row_count = 0
file_part_number = 1
total_rows_written = 0

# 增加 CSV 字段大小限制
try:
    csv.field_size_limit(1024 * 1024 * 1024)
except OverflowError:
    print("警告：无法将 csv.field_size_limit 设置为极大值，使用默认限制。长字段可能导致错误。")

def open_new_csv_file(part_num):
    """打开一个新的 CSV 文件并返回 writer 和文件对象"""
    global csv_writer, current_file, current_row_count
    if current_file:
        current_file.close()
        print(f"文件 '{current_file.name}' 写入完毕。")

    filename = f"brand_{output_prefix}{part_num}.csv"
    print(f"打开新文件: {filename}")
    # Output to the main directory
    current_file = open(filename, 'w', newline='', encoding='utf-8')
    csv_writer = csv.writer(current_file)
    csv_writer.writerow(['index', 'content']) # Header with 'index' and 'content'
    current_row_count = 0
    return csv_writer, current_file

try:
    connection = pymysql.connect(
        host=host, port=int(port), user=username, password=password,
        database=database, charset='utf8mb4'
    )
    cursor = connection.cursor() # Default cursor for fallback
    print(f"成功连接到数据库: {database}")
    # Updated print statement to reflect 10k limit
    print(f"开始从表 '{table_name}' 获取指定字段的 distinct 值并写入 CSV (每文件最多 {rows_per_file} 行)...")

    csv_writer, current_file = open_new_csv_file(file_part_number)

    for col in string_columns:
        print(f"- 正在处理字段: {col}")
        try:
            sql = f"SELECT DISTINCT `{col}` FROM `{table_name}` WHERE `{col}` IS NOT NULL AND `{col}` != ''"
            try:
                 col_cursor = connection.cursor(pymysql.cursors.SSCursor)
                 col_cursor.execute(sql)
                 print(f"  使用 SSCursor 获取字段 '{col}' 的数据...")
                 fetched_count = 0
                 for row in col_cursor:
                    if current_row_count >= rows_per_file:
                        file_part_number += 1
                        csv_writer, current_file = open_new_csv_file(file_part_number)

                    value = row[0]
                    if isinstance(value, (datetime.date, datetime.datetime)): value_str = value.isoformat()
                    elif isinstance(value, Decimal): value_str = float(value)
                    elif isinstance(value, (bytes, bytearray)): value_str = value.decode('utf-8', 'replace')
                    else: value_str = str(value)

                    csv_writer.writerow([value_str, col]) # Write value ('index') and column name ('content')
                    current_row_count += 1
                    total_rows_written += 1
                    fetched_count += 1
                 col_cursor.close()
                 print(f"  字段 '{col}' 处理完毕，获取并写入 {fetched_count} 个 distinct 值。")

            except (pymysql.ProgrammingError, pymysql.InternalError) as sse:
                 print(f"  SSCursor 不可用或出错 ({sse})，回退到默认 Cursor...")
                 cursor.execute(sql)
                 values = cursor.fetchall()
                 print(f"  使用默认 Cursor 获取到 {len(values)} 个 distinct 值。")
                 for row in values:
                     if current_row_count >= rows_per_file:
                         file_part_number += 1
                         csv_writer, current_file = open_new_csv_file(file_part_number)

                     value = row[0]
                     if isinstance(value, (datetime.date, datetime.datetime)): value_str = value.isoformat()
                     elif isinstance(value, Decimal): value_str = float(value)
                     elif isinstance(value, (bytes, bytearray)): value_str = value.decode('utf-8', 'replace')
                     else: value_str = str(value)

                     csv_writer.writerow([value_str, col]) # Write value ('index') and column name ('content')
                     current_row_count += 1
                     total_rows_written += 1
                 print(f"  字段 '{col}' (默认 Cursor) 处理完毕。")

        except pymysql.Error as e: print(f"  处理字段 '{col}' 时数据库出错: {e}")
        except Exception as e: print(f"  处理字段 '{col}' 时发生未知错误: {e}")

    print(f"\n所有字段处理完毕。总共写入 {total_rows_written} 行数据到 {file_part_number} 个 CSV 文件中。")

except pymysql.MySQLError as e: print(f"数据库连接或查询出错: {e}")
except Exception as e: print(f"发生未知错误: {e}")
finally:
    if current_file and not current_file.closed:
        current_file.close()
        print(f"文件 '{current_file.name}' 写入完毕。")
    if 'connection' in locals() and connection.open:
        if 'cursor' in locals() and cursor:
             try: cursor.close()
             except: pass
        try: connection.close()
        except: pass
        print("\n数据库连接已关闭。")
